;;; solve odes by runge-kutta-fehlberg method in stream form

(library-directories '("." "/home/cl/gits/libs-for-chezscheme"))
(library (rkf-odes)
  (export runge-kutta-fehlberg)
  (import (rnrs base)
          (rnrs control)
          (stream))

    ;; runge-kutta-fehlberg method
  (define (runge-kutta-fehlberg f init-time init-val init-step tol)
    (let str ([s (vector init-time init-val init-step)])
      (stream-cons s (str (next-state s f tol)))))

  (define (scalar-product h v)
    (vector-map (lambda (ele) (* h ele)) v))
  (define (add . vs)
    (apply vector-map (lambda eles (apply + eles)) vs))
  (define (norm v)
    (apply max (vector->list (vector-map abs v))))

  ;;key part
  (define next-state ; a state = '#(t, '#(y0,y1,..), h)
    (lambda (now-state f tol)
      (let* ([t (vector-ref now-state 0)]
             [y (vector-ref now-state 1)]
             [h (vector-ref now-state 2)]
             [k1 (scalar-product h (f t y))]
             [k2 (scalar-product h (f (+ t (/ h 4))
                                      (add y
                                           (scalar-product 1/4 k1))))]
             [k3 (scalar-product h (f (+ t (* 3/8 h))
                                      (add y
                                           (scalar-product 3/32 k1)
                                           (scalar-product 9/32 k2))))]
             [k4 (scalar-product h (f (+ t (* 12/13 h))
                                      (add y
                                           (scalar-product 1932/2197 k1)
                                           (scalar-product -7200/2197 k2)
                                           (scalar-product 7296/2197 k3))))]
             [k5 (scalar-product h (f (+ t h)
                                      (add y
                                           (scalar-product 439/216 k1)
                                           (scalar-product -8 k2)
                                           (scalar-product 3680/513 k3)
                                           (scalar-product -845/4104 k4))))]
             [k6 (scalar-product h (f (+ t (/ h 2))
                                      (add y
                                           (scalar-product -8/27 k1)
                                           (scalar-product 2 k2)
                                           (scalar-product -3544/2565 k3)
                                           (scalar-product 1859/4104 k4)
                                           (scalar-product -11/40 k5))))]
             [R (/ (norm (add (scalar-product 1/360 k1)
                              (scalar-product -128/4275 k3)
                              (scalar-product -2197/75240 k4)
                              (scalar-product 1/50 k5)
                              (scalar-product 2/55 k6)))
                   h)]
             [delta (* 0.84 (expt (/ tol R) 1/4))])
        (if (<= R tol)
            (vector (+ t h)
                    (add y
                         (scalar-product 25/216 k1)
                         (scalar-product 1408/2565 k3)
                         (scalar-product 2197/4104 k4)
                         (scalar-product -1/5 k5))
                    (if (< delta 1) h
                        (* delta h)));make step(>=h) as large as possible
            (next-state (vector (vector-ref now-state 0)
                                (vector-ref now-state 1)
                                (* delta h))
                        f tol))
        )))
  )

;;;--- ---test---
|#
(let* ([f (lambda (t y)
            (vector (vector-ref y 1)
                    (+ (* (exp (* 2 t)) (sin t))
                       (* -2 (vector-ref y 0))
                       (* 2 (vector-ref y 1)))))]
       ;;f = exp(2t) * sin(t) - 2*y0 + 2*y1
       [precise (lambda (t)
                  (* 0.2 (exp (* 2 t)) (- (sin t) (* 2 (cos t)))))]
       ;; precise solution: y0 = 0.2 * exp(2t) * (sin(t) - 2*cos(t))
       [result-stream (runge-kutta-fehlberg f 0 (vector -0.4 -0.6) 0.1 1e-2)]
       [compare (stream-map (lambda (ele)
                              (cons (precise (vector-ref ele 0))
                                    (vector-ref (vector-ref ele 1) 0)))
                            result-stream)])
  (for-each (lambda (i)
              (display (stream-ref compare i)) (newline))
            (list 0 5 10 15 20))
  )

#|


